from torch.utils.data import Dataset
import pickle

class RNFDataset(Dataset):
    def __init__(self, data_file, labels_file, transform=None):
        super(RNFDataset, self).__init__()
        self.data = self.load_data_from_file(data_file)
        self.labels = self.load_data_from_file(labels_file)
        self.transform = transform

    def load_data_from_file(self, file):
        with open(file, "rb") as f:
            return pickle.load(f)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]
        label = self.labels[idx]
        return sample, label
    

if __name__ == "__main__":
    # test
    from torchvision.transforms import ToTensor
    from torch.utils.data import DataLoader
    f_d = open("data/cifar10_resnet_features/bj_test_features.pkl", "rb")    
    f_l = open("data/cifar10_resnet_features/bj_test_label.pkl", "rb")

    data = pickle.load(f_d)
    labels = pickle.load(f_l)

    rnf_dataset = RNFDataset(data, labels, ToTensor())
    test_dataloader = DataLoader(rnf_dataset)

    for x, y in test_dataloader:
        print(x.shape)
        print(y.shape)
        print(type(x))
        break



